import numpy as np
import matplotlib.pyplot as plt


def aceleracao_gravitacional(r, G, M):
    """
    Calcula a aceleração gravitacional sobre um corpo devido a outro.

    Parâmetros:
        r (np.ndarray): vetor posição (2D).
        G (float): constante gravitacional.
        M (float): massa do corpo central.

    Retorna:
        np.ndarray: vetor aceleração (2D).
    """
    distancia = np.linalg.norm(r)
    return -G * M * r / distancia**3


def metodo_euler_cromer_orbita(r0, v0, G, M, t0, tf, dt):
    """
    Simula a órbita de um corpo usando o método de Euler-Cromer.

    Parâmetros:
        r0 (np.ndarray): vetor posição inicial (2D).
        v0 (np.ndarray): vetor velocidade inicial (2D).
        G (float): constante gravitacional.
        M (float): massa do corpo central.
        t0 (float): tempo inicial.
        tf (float): tempo final.
        dt (float): passo de tempo.

    Retorna:
        tuple: (array de tempos, array de posições, array de velocidades)
    """
    n = int((tf - t0) / dt)
    t = np.linspace(t0, tf, n)
    r = np.empty((n, 2))
    v = np.empty((n, 2))

    r[0] = r0.copy()
    v[0] = v0.copy()

    for i in range(n - 1):
        a = aceleracao_gravitacional(r[i], G, M)
        v[i + 1] = v[i] + a * dt
        r[i + 1] = r[i] + v[i + 1] * dt

    return t, r, v


def plotar_multiplas_orbitas(orbitas, r0_array, v0_array):
    """
    Plota múltiplas órbitas em um único gráfico.

    Parâmetros:
        orbitas (list): lista de tuplas (t, r, v) para cada órbita.
        r0_array (list): lista de vetores de posição inicial (2D).
        v0_array (list): lista de velocidades escalares (assumidas em y).
    """
    plt.figure()
    for i, (t, r, v) in enumerate(orbitas):
        label = f"v0 = {v0_array[i]:.2f} AU/ano"
        plt.plot(r[:, 0], r[:, 1], label=label)

    plt.xlabel("x (AU)")
    plt.ylabel("y (AU)")
    plt.title("Órbitas via Método de Euler-Cromer")
    plt.grid()
    plt.legend()
    plt.axis('equal')
    plt.show()


# ======================
# EXEMPLO DE USO
# ======================

if __name__ == "__main__":
    G = 4 * np.pi**2  # AU³ / (M * ano²)
    M = 1.0           # massa do Sol em unidades solares
    dt = 0.01
    t0 = 0.0
    tf = 2.0

    r0 = np.array([1.0, 0.0])  # Posição inicial a 1 AU
    v0_array = [5, 2 * np.pi, 7]  # Diferentes velocidades iniciais

    orbitas = []
    for v in v0_array:
        v0 = np.array([0.0, v])  # Velocidade inicial perpendicular (eixo y)
        resultado = metodo_euler_cromer_orbita(r0, v0, G, M, t0, tf, dt)
        orbitas.append(resultado)

    plotar_multiplas_orbitas(orbitas, [r0]*len(v0_array), v0_array)